# -*- coding: utf-8 -*-
"""
文件名: loc_qa_templates.py

功能:
    提供基于定位信息（归一化 bbox）生成多模态 QA 的工具函数:
      - build_box_qa(label_cn, bboxes_norm, image_path, ...) -> List[(q, a, q_type)]
    支持:
      1) 直接定位类问题（让模型指出目标位置/数量）
      2) 区域判别类问题（给定一个候选框，问是否属于该目标：正例/反例）
    反例生成策略:
      - 随机采样候选框（尽量靠近图像中心），确保与 GT 框 IoU 足够小 (默认 < 0.05)
      - 反例区域的黑像素比例不超过阈值 (默认 20%)
      - 多次尝试(默认 5 次)，失败则跳过该反例

使用样例:
    from src.utils.loc_qa_templates import build_box_qa

    qa = build_box_qa(
        label_cn="视盘",
        bboxes_norm=[[0.45,0.42,0.52,0.50], [0.60,0.60,0.66,0.68]],  # 归一化 [x1,y1,x2,y2]
        image_path="/abs/path/to/image.png",
        num_pos=2,
        num_neg=2,
    )
    for q, a, t in qa:
        print(t, q, "=>", a)

注意:
    - 输入的 bboxes_norm 为 0~1 归一化坐标 (x1 < x2, y1 < y2)
    - 本工具仅负责“生成问答对文本”，不改动图片或标注
    - 需要 PIL / numpy 依赖

作者: 你的同学
时间: 2025-10-10
"""

import os
import random
from typing import List, Tuple, Optional
from PIL import Image
import numpy as np


class BoxQAGenerator:
    """
    定位 QA 生成器:
      - 打开原图, 将归一化框转为像素框
      - 生成正例与反例问答
    """

    def __init__(
        self,
        iou_neg_threshold: float = 0.05,
        black_ratio_threshold: float = 0.20,
        neg_trials: int = 5,
        rng_seed: int = 42
    ):
        """
        参数:
            iou_neg_threshold: 反例与任一 GT 的最大 IoU 上限
            black_ratio_threshold: 候选反例框内黑像素比例阈值 (<= 该值才算可用)
            neg_trials: 生成单个反例最多尝试次数
            rng_seed: 随机种子 (确保可复现)
        """
        import random as _random
        self.iou_neg_threshold = iou_neg_threshold
        self.black_ratio_threshold = black_ratio_threshold
        self.neg_trials = neg_trials
        _random.seed(rng_seed)

        # —— 更中性的“定位”问法；明确 bbox 为 [x1,y1,x2,y2](左上-右下) —— #
        self.locate_q_tpls = [
            "请在图像中定位{}所在的区域，并给出其边界框坐标（格式为[x1,y1,x2,y2]，左上-右下）。",
            "图像中哪里是{}？请给出其边界框坐标（[x1,y1,x2,y2]，左上-右下）。",
            "请标出{}的位置，并提供对应的边界框坐标（[x1,y1,x2,y2]，左上-右下）。",
            "图像中的{}位于何处？请用边界框坐标表示（[x1,y1,x2,y2]，左上-右下）。"
        ]
        # 回答不要预设“多个/单个”的口吻；统一用“检测到 n 处”+ 坐标列表
        self.locate_a_tpls = [
            "检测到 {n} 处{label}，边界框坐标为：{boxes}。",
            "共发现 {n} 处{label}，对应边界框：{boxes}。",
            "结果：{n} 处{label}，边界框坐标如下：{boxes}。"
        ]

        # —— 判别题：只要求“是/否”，避免附加解释 —— #
        self.verify_pos_q_tpls = [
            "候选框 {} 是否属于{}？请回答是或否。",
            "给定候选框 {}，它是否对应{}？请回答是或否。",
            "下列候选区域是否为{}？候选框：{}。请回答是或否。"
        ]
        self.verify_pos_a_yes = [
            "是。",
            "是。",
            "是。"
        ]
        self.verify_neg_a_no = [
            "否。",
            "否。",
            "否。"
        ]


    # -------------- 公共主入口 --------------

    def build_box_qa(
        self,
        label_cn: str,
        bboxes_norm: List[List[float]],
        image_path: str,
        num_pos: int = 2,
        num_neg: int = 2
    ):

        if not os.path.exists(image_path):
            raise FileNotFoundError(f"原图不存在: {image_path}")
        if not bboxes_norm:
            return []

        img = Image.open(image_path).convert("RGB")
        w, h = img.size
        gts_px = [self._denorm_box(b, w, h) for b in bboxes_norm]
        qa = []

        # 1) 定位
        qa.append(self._build_locate_qa(label_cn, gts_px))

        # 2) 正例判别
        pos_src = gts_px if len(gts_px) <= num_pos else random.sample(gts_px, num_pos)
        for bx in pos_src:
            q = self._rc(self.verify_pos_q_tpls).format(self._fmt_box(bx), label_cn)
            a = self._rc(self.verify_pos_a_yes)
            qa.append((q, a, "verify"))

        # 3) 反例判别
        negs = self._sample_negatives(img, gts_px, num_neg)
        for nb in negs:
            q = self._rc(self.verify_pos_q_tpls).format(self._fmt_box(nb), label_cn)
            a = self._rc(self.verify_neg_a_no)
            qa.append((q, a, "verify"))
        return qa

    def _build_locate_qa(self, label_cn, gts_px):
        q = self._rc(self.locate_q_tpls).format(label_cn)
        boxes_str = "; ".join([self._fmt_box(b) for b in gts_px])
        # 轻量提示一次坐标规则，避免歧义
        a = self._rc(self.locate_a_tpls).format(n=len(gts_px), label=label_cn, boxes=boxes_str)
        return (q, a, "locate")
    # -------------- 内部: 反例生成 --------------

    def _sample_negatives(
        self,
        img: Image.Image,
        gts_px: List[List[int]],
        want: int
    ) -> List[List[int]]:
        """采样若干个负例框: IoU 小于阈值, 且黑像素比例 <= 阈值。"""
        if want <= 0:
            return []
        w, h = img.size
        negs: List[List[int]] = []
        attempts = 0
        max_attempts = max(want * self.neg_trials * 2, 20)
        arr = np.asarray(img)  # H, W, C

        while len(negs) < want and attempts < max_attempts:
            attempts += 1
            # 采样宽高(相对小一些), 位置尽量靠中
            box_w = int(max(8, random.gauss(mu=0.12 * w, sigma=0.05 * w)))
            box_h = int(max(8, random.gauss(mu=0.12 * h, sigma=0.05 * h)))
            cx = int(np.clip(random.gauss(mu=0.5 * w, sigma=0.15 * w), 0, w - 1))
            cy = int(np.clip(random.gauss(mu=0.5 * h, sigma=0.15 * h), 0, h - 1))

            x1 = int(np.clip(cx - box_w // 2, 0, w - 1))
            y1 = int(np.clip(cy - box_h // 2, 0, h - 1))
            x2 = int(np.clip(x1 + box_w, x1 + 1, w))
            y2 = int(np.clip(y1 + box_h, y1 + 1, h))
            cand = [x1, y1, x2, y2]

            # 与任一 GT 的 IoU 必须很小
            if any(self._iou(cand, gt) >= self.iou_neg_threshold for gt in gts_px):
                continue

            # 黑像素比例检查
            if self._black_ratio(arr, cand) > self.black_ratio_threshold:
                continue

            negs.append(cand)

        return negs

    # -------------- 内部: 工具函数 --------------

    @staticmethod
    def _denorm_box(b: List[float], w: int, h: int) -> List[int]:
        """归一化 -> 像素整型 [x1,y1,x2,y2]"""
        x1 = int(round(b[0] * w))
        y1 = int(round(b[1] * h))
        x2 = int(round(b[2] * w))
        y2 = int(round(b[3] * h))
        # 纠正顺序/越界
        x1, x2 = sorted([max(0, min(x1, w)), max(0, min(x2, w))])
        y1, y2 = sorted([max(0, min(y1, h)), max(0, min(y2, h))])
        if x2 == x1:
            x2 = min(w, x1 + 1)
        if y2 == y1:
            y2 = min(h, y1 + 1)
        return [x1, y1, x2, y2]

    @staticmethod
    def _fmt_box(b: List[int]) -> str:
        """格式化输出: (x1,y1,x2,y2)"""
        return f"({b[0]},{b[1]},{b[2]},{b[3]})"

    @staticmethod
    def _iou(a: List[int], b: List[int]) -> float:
        """两像素框 IoU"""
        ax1, ay1, ax2, ay2 = a
        bx1, by1, bx2, by2 = b
        ix1, iy1 = max(ax1, bx1), max(ay1, by1)
        ix2, iy2 = min(ax2, bx2), min(ay2, by2)
        iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1)
        inter = iw * ih
        if inter == 0:
            return 0.0
        area_a = (ax2 - ax1) * (ay2 - ay1)
        area_b = (bx2 - bx1) * (by2 - by1)
        union = area_a + area_b - inter
        return inter / union if union > 0 else 0.0

    @staticmethod
    def _black_ratio(arr: np.ndarray, box: List[int], thr: int = 8) -> float:
        """
        计算候选框内近似“黑色”像素比例
        参数:
            thr: 单通道阈值(0~255)，三个通道都 <= thr 视作黑
        """
        x1, y1, x2, y2 = box
        crop = arr[y1:y2, x1:x2, :]
        if crop.size == 0:
            return 1.0
        is_black = (crop <= thr).all(axis=2)  # HxW bool
        return float(is_black.sum()) / float(is_black.size)

    @staticmethod
    def _rc(seq: List[str]) -> str:
        """random choice 简写"""
        return random.choice(seq)


# ===== 顶层便捷函数 =====

def build_box_qa(
    gen: Optional[BoxQAGenerator],
    label_cn: str,
    bboxes_norm: List[List[float]],
    image_path: str,
    num_pos: int = 2,
    num_neg: int = 2,

) -> List[Tuple[str, str, str]]:
    """
    便捷封装: 直接用一次性实例调用

    返回:
        [(question, answer, q_type), ...]
    """
    try:
        
        return gen.build_box_qa(
            label_cn=label_cn,
            bboxes_norm=bboxes_norm,
            image_path=image_path,
            num_pos=num_pos,
            num_neg=num_neg
        )
    except Exception as e:
        # 遇到任何异常直接抛出，携带关键信息，避免静默跳过
        raise RuntimeError(
            f"[build_box_qa] 生成失败: label={label_cn}, image_path={image_path}, "
            f"num_pos={num_pos}, num_neg={num_neg}. 原因: {repr(e)}"
        ) from e
